﻿/*
 * Created by jiadong chen
 * https://jiadong-chen.medium.com/
 */

using System;
using UnityEngine;
using UnityEditor;
using System.IO;
using UnityEngine.Rendering;
using static UnityEngine.GraphicsBuffer;
using Codice.Client.BaseCommands;
using UnityEditor.PackageManager.UI;

public class AnimMapBakerWindow : EditorWindow
{
    private enum SaveStrategy
    {
        // 网格
        AnimMesh,
        // 动画贴图
        AnimMap,
        // 描边材质
        OutlineMat,
        // 基础材质
        Mat, 
        // 预置体
        Prefab 
    }

    #region FIELDS

    private const string WebBuiltInShader = "BuiltIn/WebAnimMapShader";
    private const string WebBuiltInLitShader = "BuiltIn/WebAnimMapLitShader";
    private const string WebBuiltInOutlineShader = "BuiltIn/WebAnimMapOutlineShader";

    private static readonly int MainTex = Shader.PropertyToID("_MainTex");
    private static readonly int AnimMap = Shader.PropertyToID("_AnimMap");
    private static readonly int NormalMap = Shader.PropertyToID("_NormalMap");
    private static readonly int AnimLen = Shader.PropertyToID("_AnimLen");

    private GameObject _targetGo;
    private AnimMapBaker _baker;
    private string _path = "AnimMapBaker";
    private string _subPath = "SubPath";

    private SaveStrategy _strategy = SaveStrategy.Prefab;
    private Shader _animMapShader;
    private Shader _outlineShader;

    public Mesh _targetMesh;
    public float _reduce; // 溢出校正

    private bool _isLit = false;
    private bool _isOutline = false;

    #endregion

    #region  METHODS

    [MenuItem("Tools/骨骼动画转存")]
    public static void ShowWindow()
    {
        var window = GetWindow<AnimMapBakerWindow>();
    }

    private void OnGUI()
    {
        _targetGo = (GameObject)EditorGUILayout.ObjectField("动画预置体", _targetGo, typeof(GameObject), true);
        _subPath = _targetGo == null ? _subPath : _targetGo.name;

        var resultPath = "Wrong path entered!";
        var canProceed = true;

        try
        {
            resultPath = Path.Combine(_path, _subPath);
        }
        catch (Exception e)
        {
            canProceed = false;
        }

        EditorGUILayout.LabelField(string.Format($"导出路径: {resultPath}"));
        _path = EditorGUILayout.TextField(_path);
        _subPath = EditorGUILayout.TextField(_subPath);

        _strategy = (SaveStrategy)EditorGUILayout.EnumPopup("导出类型", _strategy);

        _targetMesh = EditorGUILayout.ObjectField("原始模型", _targetMesh, typeof(Mesh), true) as Mesh;
        _reduce = EditorGUILayout.Slider("溢出校正", _reduce, 0, 10.0f);

        _isLit = EditorGUILayout.Toggle("使用光照", _isLit);
        _isOutline = EditorGUILayout.Toggle("使用描边", _isOutline);

        var prevGuiEnabled = GUI.enabled;
        GUI.enabled = canProceed;
        var requestedBake = GUILayout.Button("导出");
        GUI.enabled = prevGuiEnabled;

        // 页游烘焙
        if (requestedBake)
        {
            // 是否使用光照
            _animMapShader = Shader.Find(_isLit ? WebBuiltInLitShader : WebBuiltInShader);

            // 描边
            _outlineShader = Shader.Find(WebBuiltInOutlineShader);

            if (_targetGo == null)
            {
                EditorUtility.DisplayDialog("err", "targetGo is null！", "OK");
                return;
            }

            if (_targetMesh == null)
            {
                EditorUtility.DisplayDialog("错误", "模型不能为空", "好");
                return;
            }

            _baker = new AnimMapBaker();
            _baker.SetAnimData(_targetGo);

            var list = _baker.Bake();

            if (list == null || list.Count == 0)
            {
                EditorUtility.DisplayDialog("err", "No baked data was generated. Possible, selected wrong asset.\nYou need to select prefab with Animation component and Animations list filled with animations.", "OK");
                return;
            }

            foreach (var t in list)
            {
                var data = t;
                Save(ref data);
            }
        }
    }

    private void Save(ref BakedData data)
    {
        switch(_strategy)
        {
            case SaveStrategy.AnimMesh:
                SaveMesh(ref data);
                break;
            case SaveStrategy.AnimMap:
                SaveAnimMap(ref data);
                SaveNormalMap(ref data);
                break;
            case SaveStrategy.OutlineMat:
                SaveAsOutlineMat(ref data);
                break;
            case SaveStrategy.Mat:
                SaveAsMat(ref data);
                break;
            case SaveStrategy.Prefab:
                SaveAsPrefab(ref data);
                break;
        }
        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();
    }

    private Mesh SaveMesh(ref BakedData data)
    {
        if (_targetMesh == null)
        {
            EditorUtility.DisplayDialog("错误", "模型不能为空", "好");
        }

        var folderPath = CreateFolder();
        var path = Path.Combine(folderPath, $"{_targetMesh.name}.asset");
        Mesh assetMesh = AssetDatabase.LoadAssetAtPath<Mesh>(path);

        if (assetMesh != null)
        {
            return assetMesh;
        }

        Color[] colors = new Color[_targetMesh.vertexCount];
        for (var j = 0; j < _targetMesh.vertexCount; j++)
        {
            colors[j] = new Color(j, 0, 0, 0);
        }

        Mesh instanceMesh = Instantiate(_targetMesh);
        instanceMesh.colors = colors;

        AssetDatabase.CreateAsset(instanceMesh, path);

        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();

        return instanceMesh;
    }

    private Texture2D SaveNormalMap(ref BakedData data)
    {
        var folderPath = CreateFolder();
        var normalMap = new Texture2D(data.AnimMapWidth, data.AnimMapHeight, TextureFormat.RGBAHalf, false);
        normalMap.LoadRawTextureData(data.RawNormalMap);

        var bytes = normalMap.EncodeToPNG();
        string path = Path.Combine(folderPath, data.Name + ".normalMap.png");
        File.WriteAllBytes(path, bytes);

        // Import Resources
        AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceUpdate);

        // 设置贴图参数
        TextureImporter textureIm = AssetImporter.GetAtPath(path) as TextureImporter;
        textureIm.isReadable = false;
        textureIm.anisoLevel = 9;
        textureIm.mipmapEnabled = false;
        textureIm.wrapMode = TextureWrapMode.Repeat;
        textureIm.maxTextureSize = normalMap.width;
        textureIm.sRGBTexture = false;

        TextureImporterPlatformSettings settings = textureIm.GetDefaultPlatformTextureSettings();
        settings.resizeAlgorithm = TextureResizeAlgorithm.Mitchell;
        settings.format = TextureImporterFormat.RGBA32;
        textureIm.SetPlatformTextureSettings(settings);

        AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceUpdate); // Reflash

        normalMap = AssetDatabase.LoadAssetAtPath<Texture2D>(path);

        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();

        return normalMap;
    }

    private Texture2D SaveAnimMap(ref BakedData data)
    {
        var folderPath = CreateFolder();
        var animMap = new Texture2D(data.AnimMapWidth, data.AnimMapHeight, TextureFormat.RGBAHalf, false);
        animMap.LoadRawTextureData(data.RawAnimMap);

        Vector3 maxVector = new Vector3(float.MinValue, float.MinValue, float.MinValue);
        Vector3 minVector = new Vector3(float.MaxValue, float.MaxValue, float.MaxValue);

        for (var i = 0; i < _targetMesh.vertexCount; i++)
        {
            var vertex = _targetMesh.vertices[i];

            // 记录最大
            if (vertex.x > maxVector.x)
            {
                maxVector.x = vertex.x;
            }

            if (vertex.y > maxVector.y)
            {
                maxVector.y = vertex.y;
            }

            if (vertex.z > maxVector.z)
            {
                maxVector.z = vertex.z;
            }

            // 记录最小
            if (vertex.x < minVector.x)
            {
                minVector.x = vertex.x;
            }

            if (vertex.y < minVector.y)
            {
                minVector.y = vertex.y;
            }

            if (vertex.z < minVector.z)
            {
                minVector.z = vertex.z;
            }
        }

        // 模型空间最大距离长度
        float maxLength = 0;

        if (maxLength < maxVector.x - minVector.x)
        {
            maxLength = maxVector.x - minVector.x;
        }

        if (maxLength < maxVector.y - minVector.y)
        {
            maxLength = maxVector.y - minVector.y;
        }

        if (maxLength < maxVector.z - minVector.z)
        {
            maxLength = maxVector.z - minVector.z;
        }

        float reduceLenght = maxLength * _reduce;

        minVector.x = minVector.x - reduceLenght * 0.5f;
        minVector.y = minVector.y - reduceLenght * 0.5f;
        minVector.z = minVector.z - reduceLenght * 0.5f;

        maxLength = maxLength + reduceLenght;

        // 把坐标变换到第一象限并且限制在0到1内
        Color[] mapColors = animMap.GetPixels();
        for (var i = 0; i < mapColors.Length; i++)
        {
            // 坐标归零
            mapColors[i].r = (mapColors[i].r - minVector.x) / maxLength;
            mapColors[i].g = (mapColors[i].g - minVector.y) / maxLength;
            mapColors[i].b = (mapColors[i].b - minVector.z) / maxLength;
            mapColors[i].a = 1;
        }
        animMap.SetPixels(mapColors);

        var bytes = animMap.EncodeToPNG();
        string path = Path.Combine(folderPath, data.Name + ".animMap.png");
        File.WriteAllBytes(path, bytes);

        // Import Resources
        AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceUpdate);

        // 设置贴图参数
        TextureImporter textureIm = AssetImporter.GetAtPath(path) as TextureImporter;
        textureIm.isReadable = false;
        textureIm.anisoLevel = 9;
        textureIm.mipmapEnabled = false;
        textureIm.wrapMode = TextureWrapMode.Repeat;
        textureIm.maxTextureSize = animMap.width;
        textureIm.sRGBTexture = false;

        TextureImporterPlatformSettings settings = textureIm.GetDefaultPlatformTextureSettings();
        settings.resizeAlgorithm = TextureResizeAlgorithm.Mitchell;
        settings.format = TextureImporterFormat.RGBA32;
        textureIm.SetPlatformTextureSettings(settings);

        AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceUpdate); // Reflash

        animMap = AssetDatabase.LoadAssetAtPath<Texture2D>(path);

        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();

        return animMap;
    }

    private Material SaveAsMat(ref BakedData data)
    {
        if(_animMapShader == null)
        {
            EditorUtility.DisplayDialog("err", "shader is null!!", "OK");
            return null;
        }

        if(_targetGo == null || !_targetGo.GetComponentInChildren<SkinnedMeshRenderer>())
        {
            EditorUtility.DisplayDialog("err", "SkinnedMeshRender is null!!", "OK");
            return null;
        }

        var smr = _targetGo.GetComponentInChildren<SkinnedMeshRenderer>();
        var mat = new Material(_animMapShader);
        var animMap = SaveAnimMap(ref data);

        mat.SetTexture(MainTex, smr.sharedMaterial.mainTexture);
        mat.SetTexture(AnimMap, animMap);
        mat.SetFloat(AnimLen, data.AnimLen);

        if (_isLit)
        {
            var normalMap = SaveNormalMap(ref data);
            mat.SetTexture(NormalMap, normalMap);
        }

        var folderPath = CreateFolder();
        AssetDatabase.CreateAsset(mat, Path.Combine(folderPath, $"{data.Name}.mat"));

        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();

        return mat;
    }

    private Material SaveAsOutlineMat(ref BakedData data)
    {
        if (_outlineShader == null)
        {
            EditorUtility.DisplayDialog("err", "shader is null!!", "OK");
            return null;
        }

        if (_targetGo == null || !_targetGo.GetComponentInChildren<SkinnedMeshRenderer>())
        {
            EditorUtility.DisplayDialog("err", "SkinnedMeshRender is null!!", "OK");
            return null;
        }

        var smr = _targetGo.GetComponentInChildren<SkinnedMeshRenderer>();
        var mat = new Material(_outlineShader);
        var animMap = SaveAnimMap(ref data);
        var normalMap = SaveNormalMap(ref data);

        mat.SetTexture(AnimMap, animMap);
        mat.SetTexture(NormalMap, normalMap);
        mat.SetFloat(AnimLen, data.AnimLen);

        var folderPath = CreateFolder();
        AssetDatabase.CreateAsset(mat, Path.Combine(folderPath, $"{data.Name}.outline.mat"));

        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();

        return mat;
    }

    private void SaveAsPrefab(ref BakedData data)
    {
        var mat = SaveAsMat(ref data);

        // 基础材质
        if (mat == null)
        {
            EditorUtility.DisplayDialog("err", "基础材质 is null!!", "OK");
            return;
        }

        var go = new GameObject();

        if (_isOutline)
        {
            var outlineMat = SaveAsOutlineMat(ref data);

            // 描边材质
            if (outlineMat == null)
            {
                EditorUtility.DisplayDialog("err", "描边材质 is null!!", "OK");
                return;
            }

            Material[] sharedMaterials = { mat, outlineMat };
            go.AddComponent<MeshRenderer>().sharedMaterials = sharedMaterials;
        }
        else
        {
            go.AddComponent<MeshRenderer>().sharedMaterial = mat;
        }

        var mesh = SaveMesh(ref data);

        if (mesh == null)
        {
            EditorUtility.DisplayDialog("错误", "网格为空", "好");
            return;
        }

        go.AddComponent<MeshFilter>().sharedMesh = mesh;

        var folderPath = CreateFolder();
        PrefabUtility.SaveAsPrefabAsset(go, Path.Combine(folderPath, $"{data.Name}.prefab").Replace("\\", "/"));
    }

    private string CreateFolder()
    {
        var folderPath = Path.Combine("Assets", _path);

        if (!AssetDatabase.IsValidFolder(folderPath))
        {
            AssetDatabase.CreateFolder("Assets", _path);
        }

        folderPath = Path.Combine("Assets/" + _path, _subPath);

        if (!AssetDatabase.IsValidFolder(folderPath))
        {
            AssetDatabase.CreateFolder("Assets/" + _path, _subPath);
        }
        return folderPath;
    }

    #endregion

}
